updated: pipeline

app.py

@@ -4,10 +4,10 @@ from PIL import Image
 from explain import get_results, reproduce
 
 
-def classify_and_explain(image, object_detection=False):
+def classify_and_explain(image, object_detection=False, remove_bg=False):
     reproduce()
     # This function will classify the image and return a list of image paths
-    list_of_images = get_results(img_for_testing=image, od=object_detection)
+    list_of_images = get_results(img_for_testing=image, od=object_detection, remove_bg=remove_bg)
     return list_of_images
 
 
@@ -23,8 +23,8 @@ def get_examples():
     ]
 
     return [
-        [Image.open(i), True] for i in glob("samples/*") if i not in od_off_examples
-    ] + [[Image.open(i), False] for i in glob("samples/*") if i in od_off_examples]
+        [Image.open(i), True, False] for i in glob("samples/*") if i not in od_off_examples
+    ] + [[Image.open(i), False, False] for i in glob("samples/*") if i in od_off_examples]
 
 
 demo = gr.Interface(
@@ -34,6 +34,9 @@ demo = gr.Interface(
         gr.Checkbox(
             label="Extract Leaves", info="What to extract leafs before classification"
         ),
+        gr.Checkbox(
+            label="Remove Background", info="Remove the background and make it white"
+        ),
     ],
     outputs="gallery",
     examples=get_examples(),

explain.py

@@ -16,7 +16,7 @@ from pytorch_grad_cam.utils.image import show_cam_on_image, deprocess_image
 
 from ultralytics import YOLO
 
-# from rembg import remove
+from rembg import remove
 import uuid
 
 
@@ -25,7 +25,7 @@ model_path = "efficientnet-b2.pth"
 model_name = "efficientnet_b2"
 YOLO_MODEL_WEIGHTS = "yolo-v11-best.pt"
 classes = ["Healthy", "Resistant", "Susceptible"]
-resizing_transforms = transforms.Compose([transforms.CenterCrop(224)])
+resizing_transforms = transforms.Compose([transforms.CenterCrop(256)])
 
 
 # Function definitions
@@ -141,7 +141,7 @@ gbp_model = GuidedBackpropReLUModel(model=model, device="cpu")
 yolo_model = YOLO(YOLO_MODEL_WEIGHTS)
 
 
-def get_results(img_path=None, img_for_testing=None, od=False):
+def get_results(img_path=None, img_for_testing=None, od=False, remove_bg=False):
     if img_path is None and img_for_testing is None:
         raise ValueError("Either img_path or img_for_testing should be provided.")
 
@@ -161,14 +161,15 @@ def get_results(img_path=None, img_for_testing=None, od=False):
             bbox = result.boxes.xyxy[0].cpu().numpy().astype(int)
             bbox_image = image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
 
-            # bbox_image = remove(bbox_image).convert("RGB")
-            # bbox_image = Image.fromarray(
-            #     np.where(
-            #         np.array(bbox_image) == [0, 0, 0],
-            #         [255, 255, 255],
-            #         np.array(bbox_image),
-            #     ).astype(np.uint8)
-            # )
+            if remove_bg:
+                bbox_image = remove(bbox_image).convert("RGB")
+                bbox_image = Image.fromarray(
+                    np.where(
+                        np.array(bbox_image) == [0, 0, 0],
+                        [255, 255, 255],
+                        np.array(bbox_image),
+                    ).astype(np.uint8)
+                )
 
             res = make_prediction_and_explain(bbox_image)
             save_explanation_results(res, save_path)

requirements.txt

@@ -8,4 +8,5 @@ wandb
 seaborn
 matplotlib
 ultralytics
-grad-cam
+grad-cam
+rembg[cpu]