Update app.py

app.py

@@ -75,12 +75,12 @@ def detect_objects(model_name,url_input,image_input,threshold):
     
     return viz_img   
 
-def detect_objects2(model_name,url_input,image_input,threshold):
+def detect_objects2(model_name,url_input,image_input,threshold,type2):
     
     #Extract model and feature extractor
     feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
     
-    
+    xxresult=0
         
     model = DetrForObjectDetection.from_pretrained(model_name)
         
@@ -97,8 +97,33 @@ def detect_objects2(model_name,url_input,image_input,threshold):
     viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
 
     keep = processed_outputs["scores"] > threshold
+    det_lab = processed_outputs["labels"][keep].tolist() 
+    det_lab.count(1)
+    if det_lab.count(1) > 0:
+        total_text="Trench is Detected \n"
+    else:
+        total_text="Trench is NOT Detected \n"
+        xxresult=1
+        
     
-    return processed_outputs["labels"][keep].tolist() 
+    if det_lab.count(4) > 0:
+        total_text+="Measuring Tape (Vertical) for measuring Depth is Detected \n"
+    else:
+        total_text+="Measuring Tape (Vertical) for measuring Depth is NOT Detected \n
+        if type2=="Trench Depth Measurement":
+            xxresult=1
+
+    if det_lab.count(5) > 0:
+        total_text+="Measuring Tape (Horizontal) for measuring Width is Detected \n"
+    else:
+        total_text+="Measuring Tape (Horizontal) for measuring Width is NOT Detected \n"
+        if type2=="Trench Width Measurement":
+            xxresult=1
+    if xxresult==0:
+        text2 = "The photo is ACCEPTED"
+    else:
+        text2 = "The photo is NOT ACCEPTED"
+    return total_text, text2
     
 def set_example_image(example: list) -> dict:
     return gr.Image.update(value=example[0])
@@ -118,7 +143,7 @@ Links to HuggingFace Models:
 """
 
 models = ["omarhkh/detr-finetuned-omar8"]
-
+types_class = ["Trench Depth Measurement", "Trench Width Measurement"]
 
 css = '''
 h1#title {
@@ -135,6 +160,7 @@ with demo:
     
     
     options = gr.Dropdown(choices=models,label='Select Object Detection Model',show_label=True)
+    options2 = gr.Dropdown(choices=types_class,label='Select Classification Type',show_label=True)
     slider_input = gr.Slider(minimum=0.1,maximum=1,value=0.7,label='Prediction Threshold')
     
     with gr.Tabs():
@@ -151,10 +177,12 @@ with demo:
             img_but = gr.Button('Detect')
 
         with gr.Blocks():
-            name = gr.Textbox(label="Name")
-            output = gr.Textbox(label="Results")
+            name = gr.Textbox(label="Final Result")
+            output = gr.Textbox(label="Reason for the results")
             greet_btn = gr.Button("Results")
-            greet_btn.click(fn=detect_objects2, inputs=[options,img_input,img_input,slider_input], outputs=output, queue=True)
+            greet_btn.click(fn=detect_objects2[0], inputs=[options,img_input,img_input,slider_input,options2], outputs=output, queue=True)
+            name.change(fn=detect_objects2[1], inputs=[options,img_input,img_input,slider_input,options2], outputs=name, queue=True)
+