Refactor : 오류 위험이 있는 코드 전부 제거

app.py

@@ -131,26 +131,24 @@ def sepia(inputs, button_text):
     #     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
+    inputs_segmentation = feature_extractor(images=input_img, return_tensors="tf")
+    outputs_segmentation = model_segmentation(**inputs_segmentation)
+    logits_segmentation = outputs_segmentation.logits
 
-        logits_segmentation = tf.transpose(logits_segmentation, [0, 2, 3, 1])
-        logits_segmentation = tf.image.resize(logits_segmentation, input_img.size[::-1])
-        seg = tf.math.argmax(logits_segmentation, axis=-1)[0]
+    logits_segmentation = tf.transpose(logits_segmentation, [0, 2, 3, 1])
+    logits_segmentation = tf.image.resize(logits_segmentation, input_img.size[::-1])
+    seg = tf.math.argmax(logits_segmentation, axis=-1)[0]
 
-        color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
-        for label, color in enumerate(colormap):
-            color_seg[seg.numpy() == label, :] = color
+    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
+    for label, color in enumerate(colormap):
+        color_seg[seg.numpy() == label, :] = color
 
-        pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
-        pred_img = pred_img.astype(np.uint8)
+    pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
+    pred_img = pred_img.astype(np.uint8)
+
+    fig = draw_plot(pred_img, seg)
+    return fig
 
-        fig = draw_plot(pred_img, seg)
-        return fig
-    
-    return "Please select 'detection' or 'segmentation'."
 
 def on_button_click(inputs):
     """버튼 클릭 이벤트 핸들러"""
@@ -166,18 +164,15 @@ dropdown = gr.Dropdown(
             ["segmentation"], label="Menu", info="Chose Segmentation!"
         )
 
-demo = gr.Interface(
-    fn=sepia,
-    inputs=[gr.Image(shape=(400, 600))],
-    outputs=["plot"],
-    examples=[
-        ["01.jpg", "1"],
-        ["02.jpeg", "2"],
-        ["03.jpeg", "3"],
-        ["04.jpeg", "4"],
-    ],
-    allow_flagging="never",
-)
-
+demo = gr.Interface(fn=sepia,
+                    inputs=[gr.Image(shape=(400, 600)), dropdown],
+                    outputs=["plot"],
+                    examples= [
+                        ["01.jpg", "1"], 
+                        ["02.jpeg", "2"],
+                        ["03.jpeg", "3"],
+                        ["04.jpeg", "4"],
+                    ],
+                    allow_flagging="never",)
 
 demo.launch()