add: OD option

app.py

@@ -4,18 +4,38 @@ from PIL import Image
 from explain import get_results, reproduce
 
 
-def classify_and_explain(image):
+def classify_and_explain(image, object_detection=False):
     reproduce()
     # This function will classify the image and return a list of image paths
-    list_of_images = get_results(img_for_testing=image)
+    list_of_images = get_results(img_for_testing=image, od=object_detection)
     return list_of_images
 
 
 def get_examples():
-    return [Image.open(i) for i in glob("samples/*")]
+
+    od_off_examples = [
+        "samples/DSC_0315.jpg",
+        "samples/20210401_123624.jpg",
+        "samples/IMG_1299.jpg",
+        "samples/20210420_112400.jpg",
+        "samples/IMG_1300.jpg",
+        "samples/20210420_112406.jpg",
+    ]
+
+    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]
 
 
 demo = gr.Interface(
-    fn=classify_and_explain, inputs="image", outputs="gallery", examples=get_examples()
+    fn=classify_and_explain,
+    inputs=[
+        "image",
+        gr.Checkbox(
+            label="Extract Leaves", info="What to extract leafs before classification"
+        ),
+    ],
+    outputs="gallery",
+    examples=get_examples(),
 )
 demo.launch()

explain.py

@@ -15,6 +15,7 @@ from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 from pytorch_grad_cam.utils.image import show_cam_on_image, deprocess_image
 
 from ultralytics import YOLO
+
 # from rembg import remove
 import uuid
 
@@ -132,7 +133,7 @@ def save_explanation_results(res, path):
 
 
 model, image_transform = get_model(model_name)
-model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
 model.train()
 target_layers = [model.conv_head]
 gbp_model = GuidedBackpropReLUModel(model=model, device="cpu")
@@ -140,38 +141,44 @@ gbp_model = GuidedBackpropReLUModel(model=model, device="cpu")
 yolo_model = YOLO(YOLO_MODEL_WEIGHTS)
 
 
-def get_results(img_path=None, img_for_testing=None):
+def get_results(img_path=None, img_for_testing=None, od=False):
     if img_path is None and img_for_testing is None:
         raise ValueError("Either img_path or img_for_testing should be provided.")
 
     if img_path is not None:
-        results = yolo_model(img_path)
         image = Image.open(img_path)
 
     if img_for_testing is not None:
-        results = yolo_model(img_for_testing)
         image = Image.fromarray(img_for_testing)
 
     result_paths = []
 
-    for i, result in enumerate(results):
+    if od:
+        results = yolo_model(img_path if img_path else img_for_testing)
+        for i, result in enumerate(results):
+            unique_id = uuid.uuid4().hex
+            save_path = f"/tmp/with-bg-result-{unique_id}.png"
+            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)
+            # )
+
+            res = make_prediction_and_explain(bbox_image)
+            save_explanation_results(res, save_path)
+
+            result_paths.append(save_path)
+    else:
         unique_id = uuid.uuid4().hex
-        save_path = f"/tmp/with-white-bg-result-{unique_id}.png"
-        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)
-        # )
-
-        res = make_prediction_and_explain(bbox_image)
+        save_path = f"/tmp/with-bg-result-{unique_id}.png"
+        res = make_prediction_and_explain(image)
         save_explanation_results(res, save_path)
-
         result_paths.append(save_path)
 
     return result_paths