import torch import timm from collections import OrderedDict import types from PIL import Image import os import re from torchvision import transforms from concurrent.futures import ThreadPoolExecutor, as_completed import io def preprocess_image(image): # simulate WhatsApp compression buffer = io.BytesIO() image.save(buffer, format="JPEG", quality=75) # 🔥 KEY buffer.seek(0) image = Image.open(buffer).convert("RGB") # final resize image = image.resize((224, 224)) return image # Paths to trained models MODEL_PATHS = { "obstacle": "/var/www/html/ai-image-ml/obstacle_classifier.pt" , "moisture":"/var/www/html/ai-image-ml/moisture_classifier.pt", # "soil": "/var/www/html/agri_image_ai/soil_visibility_classifier.pt", "ruler": "/var/www/html/ai-image-ml/ruler_classifier.pt", "clarity": "/var/www/html/ai-image-ml/clarity_classifier.pt", "pipe":"/var/www/html/ai-image-ml/pipe_classifier.pt" } ######################################################## import os import torch from collections import OrderedDict import timm import types from torchvision.models import resnet # for allowlisting if needed # your MODEL_ARCH and MODEL_PATHS as before MODEL_ARCH = { "obstacle": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1), "clarity": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1), "pipe": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1), "ruler": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1), "moisture": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1) } def strip_module_prefix(state_dict): new_state = OrderedDict() for k, v in state_dict.items(): new_key = k[len("module."):] if k.startswith("module.") else k new_state[new_key] = v return new_state def try_extract_state_dict(checkpoint): # Returns (state_dict_or_None, model_obj_or_None) if isinstance(checkpoint, OrderedDict) or isinstance(checkpoint, dict): # if it's a plain mapping with string keys, assume state_dict if all(isinstance(k, str) for k in checkpoint.keys()): return strip_module_prefix(checkpoint), None if isinstance(checkpoint, dict): for key in ("state_dict", "model_state_dict", "model"): if key in checkpoint: candidate = checkpoint[key] if isinstance(candidate, (OrderedDict, dict)): return strip_module_prefix(candidate), None if hasattr(candidate, "eval"): return None, candidate # if it's an nn.Module-like object if hasattr(checkpoint, "eval") and isinstance(checkpoint.eval, types.MethodType): return None, checkpoint return None, None def robust_torch_load(path): """ Attempt to load checkpoint safely: 1) try weights_only=True (safe) 2) on specific failure, retry inside safe_globals for known classes with weights_only=False Returns whatever torch.load returns (OrderedDict, dict, or nn.Module) """ try: # try the safe weights-only load first (no unpickle execution) return torch.load(path, map_location=torch.device("cpu"), weights_only=True) except Exception as e: msg = str(e) # If message indicates weights-only load failed and mentions unsupported global, try allowlist if "Weights only load failed" in msg or "Unsupported global" in msg or "UnsupportedGlobal" in msg: # --- ONLY do this for files you trust --- try: # Temporarily allowlist the ResNet global for unpickling with torch.serialization.safe_globals([resnet.ResNet]): return torch.load(path, map_location=torch.device("cpu"), weights_only=False) except Exception as e2: # re-raise the second failure for clarity raise RuntimeError(f"Failed loading checkpoint even inside safe_globals: {e2}") from e2 # otherwise re-raise original error raise # MAIN loader loop - creates architecture then loads the checkpoint (weights or full model) models_dict = {} for key, path in MODEL_PATHS.items(): try: if not os.path.exists(path): raise FileNotFoundError(f"path not found: {path}") if key not in MODEL_ARCH: raise KeyError(f"No architecture defined for key '{key}'") # create architecture (we'll load weights into this if the checkpoint is state_dict-like) arch_model = MODEL_ARCH[key]() # robustly load the checkpoint (returns dict/OrderedDict or nn.Module) checkpoint = robust_torch_load(path) state_dict, loaded_model = try_extract_state_dict(checkpoint) if state_dict is not None: # load into architecture arch_model.load_state_dict(state_dict, strict=False) arch_model.eval() models_dict[key] = arch_model print(f"✅ Loaded weights into architecture for: {key}") elif loaded_model is not None: # checkpoint was a saved model object; unwrap if DataParallel real_model = getattr(loaded_model, "module", loaded_model) real_model.eval() models_dict[key] = real_model print(f"✅ Loaded full saved model for: {key}") else: raise RuntimeError("Unrecognized checkpoint format (neither state_dict nor model object).") except Exception as e: print(f"❌ Failed to load model '{key}' from {path}: {e}") raise ######################################################## # # Define model architecture per classifier # MODEL_ARCH = { # "obstacle": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1) # , # # "soil": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1), # # "ruler": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1), # "clarity": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1), # "pipe": lambda: timm.create_model('efficientnet_b0', pretrained=False, num_classes=1) # } # # Load models # models_dict = {} # for key, path in MODEL_PATHS.items(): # try: # model = MODEL_ARCH[key]() # torch.serialization.add_safe_globals([resnet.ResNet]) # #model.load_state_dict(torch.load(path, map_location=torch.device("cpu"))) # model = torch.load(path, map_location=torch.device("cpu")) # model.eval() # models_dict[key] = model # print(f"✅ Loaded model: {key}") # except Exception as e: # print(f"❌ Failed to load model '{key}' from {path}: {e}") # raise # Image transform transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # Prediction for a single image def predict_all(image_path,activity): print("activity ,", activity) try: image = Image.open(image_path).convert("RGB") image = preprocess_image(image) if image.size[0] < 50 or image.size[1] < 50: return {"accepted": False, "reasons": ["Image too small."]} input_tensor = transform(image).unsqueeze(0) results = {} with torch.no_grad(): for key, model in models_dict.items(): output = model(input_tensor) prob = torch.sigmoid(output).item() results[key] = prob > 0.5 failed = [] # if not results["soil"]: # failed.append("Soil not visible inside pipe.") # if not results["ruler"]: # failed.append("Ruler not present inside pipe.") if not results["clarity"]: failed.append("Image is unclear.") if results["obstacle"]: failed.append("Pipe is obstructed.") if not results["pipe"]: failed.append("No pipe in the Image.") if re.search("select", activity, re.IGNORECASE): if results["moisture"]: failed.append("Soil has moisture.") # if results["ruler"]: # failed.append("Ruler not present inside pipe.") # if results["ruler"] and not results["soil"]: # failed.append("Ruler is present but soil is not visible through pipe.") return { "accepted": len(failed) == 0, "reasons": failed } except Exception as e: return {"accepted": False, "reasons": [f"Error: {str(e)}"]} # Parallel batch filtering # def batch_filter_images_parallel(source_dir, accepted_dir, max_workers=4): # os.makedirs(accepted_dir, exist_ok=True) # results = [] # def process_image(filename): # path = os.path.join(source_dir, filename) # result = predict_all(path) # result["filename"] = filename # if result["accepted"]: # os.rename(path, os.path.join(accepted_dir, filename)) # return result # image_files = [ # f for f in os.listdir(source_dir) # if f.lower().endswith((".jpg", ".jpeg", ".png")) # ] # with ThreadPoolExecutor(max_workers=max_workers) as executor: # future_to_file = {executor.submit(process_image, f): f for f in image_files} # for future in as_completed(future_to_file): # try: # result = future.result() # results.append(result) # status = "✅ Accepted" if result["accepted"] else "❌ Rejected" # print(f"{status}: {result['filename']} → {', '.join(result['reasons']) if result['reasons'] else 'No issues'}") # except Exception as e: # filename = future_to_file[future] # print(f"❌ Error processing {filename}: {str(e)}") # results.append({ # "filename": filename, # "accepted": False, # "reasons": [f"Error: {str(e)}"] # }) # return results ###################################################### def batch_filter_images_parallel(source_dir, accepted_dir, max_workers=4): os.makedirs(accepted_dir, exist_ok=True) results = [] any_accepted = False # track if any image passed def process_image(filename): path = os.path.join(source_dir, filename) result = predict_all(path) result["filename"] = filename if result["accepted"]: os.rename(path, os.path.join(accepted_dir, filename)) return result image_files = [ f for f in os.listdir(source_dir) if f.lower().endswith((".jpg", ".jpeg", ".png")) ] with ThreadPoolExecutor(max_workers=max_workers) as executor: future_to_file = {executor.submit(process_image, f): f for f in image_files} for future in as_completed(future_to_file): try: result = future.result() results.append(result) if result["accepted"]: any_accepted = True # ⭐ track success status = "✅ Accepted" if result["accepted"] else "❌ Rejected" print(f"{status}: {result['filename']} → {', '.join(result['reasons']) if result['reasons'] else 'No issues'}") except Exception as e: filename = future_to_file[future] print(f"❌ Error processing {filename}: {str(e)}") results.append({ "filename": filename, "accepted": False, "reasons": [f"Error: {str(e)}"] }) # ⭐ If any image accepted → mark all as accepted if any_accepted: for r in results: r["accepted"] = True r["reasons"] = [] return results ######################################################## # Run batch if executed directly if __name__ == "__main__": SOURCE_DIR = "raw_images" ACCEPTED_DIR = "accepted_images" results = batch_filter_images_parallel(SOURCE_DIR, ACCEPTED_DIR)