Refactor : detection 코드 주석

app.py

@@ -9,6 +9,8 @@ import torch
 import tensorflow as tf
 from PIL import ImageDraw
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # image segmentation 모델
 feature_extractor = SegformerFeatureExtractor.from_pretrained(
     "nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
@@ -18,12 +20,12 @@ model_segmentation = TFSegformerForSemanticSegmentation.from_pretrained(
 )
 
 # image detection 모델
-processor_detection = DetrImageProcessor.from_pretrained(
-    "facebook/detr-resnet-50"
-)
-model_detection = DetrForObjectDetection.from_pretrained(
-    "facebook/detr-resnet-50"
-)
+# processor_detection = DetrImageProcessor.from_pretrained(
+#     "facebook/detr-resnet-50", revision="no_timm"
+# )
+# model_detection = DetrForObjectDetection.from_pretrained(
+#     "facebook/detr-resnet-50", revision="no_timm"
+# )
 
 
 def ade_palette():
@@ -100,36 +102,36 @@ def sepia(inputs, button_text):
     """객체 검출 또는 세그멘테이션을 수행하고 결과를 반환하는 함수입니다."""
 
     input_img = Image.fromarray(inputs)
-    if button_text == "detection":
-        inputs_detection = processor_detection(images=input_img, return_tensors="pt")
-        outputs_detection = model_detection(**inputs_detection)
-
-        target_sizes = torch.tensor([input_img.size[::-1]])
-        results_detection = processor_detection.post_process_object_detection(
-            outputs_detection, target_sizes=target_sizes, threshold=0.9
-        )[0]
-
-        draw = ImageDraw.Draw(input_img)
-        for score, label, box in zip(
-            results_detection["scores"],
-            results_detection["labels"],
-            results_detection["boxes"],
-        ):
-            box = [round(i, 2) for i in box.tolist()]
-            label_name = model_detection.config.id2label[label.item()]
-            print(
-                f"Detected {label_name} with confidence "
-                f"{round(score.item(), 3)} at location {box}"
-            )
-            draw.rectangle(box, outline="red", width=3)
-            draw.text((box[0], box[1]), label_name, fill="red", font=None)
-
-        fig = plt.figure(figsize=(20, 15))
-        plt.imshow(input_img)
-        plt.axis("off")
-        return fig
-
-    elif button_text == "segmentation":
+    # if button_text == "detection":
+    #     inputs_detection = processor_detection(images=input_img, return_tensors="pt")
+    #     outputs_detection = model_detection(**inputs_detection)
+
+    #     target_sizes = torch.tensor([input_img.size[::-1]])
+    #     results_detection = processor_detection.post_process_object_detection(
+    #         outputs_detection, target_sizes=target_sizes, threshold=0.9
+    #     )[0]
+
+    #     draw = ImageDraw.Draw(input_img)
+    #     for score, label, box in zip(
+    #         results_detection["scores"],
+    #         results_detection["labels"],
+    #         results_detection["boxes"],
+    #     ):
+    #         box = [round(i, 2) for i in box.tolist()]
+    #         label_name = model_detection.config.id2label[label.item()]
+    #         print(
+    #             f"Detected {label_name} with confidence "
+    #             f"{round(score.item(), 3)} at location {box}"
+    #         )
+    #         draw.rectangle(box, outline="red", width=3)
+    #         draw.text((box[0], box[1]), label_name, fill="red", font=None)
+
+    #     fig = plt.figure(figsize=(20, 15))
+    #     plt.imshow(input_img)
+    #     plt.axis("off")
+    #     return fig
+
+    if button_text == "segmentation":
         inputs_segmentation = feature_extractor(images=input_img, return_tensors="tf")
         outputs_segmentation = model_segmentation(**inputs_segmentation)
         logits_segmentation = outputs_segmentation.logits
@@ -155,13 +157,13 @@ def on_button_click(inputs):
     image_path, selected_option = inputs
     if selected_option == "dropout":
         # 'dropout'이면 두 가지 중에 하나를 랜덤으로 선택
-        selected_option = np.random.choice(["detection", "segmentation"])
+        selected_option = np.random.choice(["segmentation"])
     
     return sepia(image_path, selected_option)
 
 # Gr.Dropdown을 사용하여 옵션을 선택할 수 있도록 변경
 dropdown = gr.Dropdown(
-            ["detection", "segmentation"], label="Menu", info="Select One!"
+            ["segmentation"], label="Menu", info="Chose Segmentation!"
         )
 
 demo = gr.Interface(
@@ -177,4 +179,4 @@ demo = gr.Interface(
     allow_flagging="never",
 )
 
-demo.launch(share=False)
+demo.launch()