app

app.py

@@ -0,0 +1,85 @@
+import gradio as gr
+from PIL import Image,ImageDraw, ImageFont
+import sys
+import os
+import torch
+from util import Detection
+
+face_model = os.environ.get('FACE_MODEL')
+age_model = os.environ.get('AGE_MODEL')
+
+torch.hub.download_url_to_file(face_model, 'face_model.pt')
+torch.hub.download_url_to_file(age_model, 'age_model.pt')
+
+sys.path.append("./")
+sys.path.append("./yolov5")
+
+age_model_ts = torch.jit.load("age_model.pt")
+
+from yolov5.detect import predict, load_yolo_model
+
+# Model
+
+model, stride, names, pt, jit, onnx, engine = load_yolo_model("face_model.pt", imgsz=[320,320])
+
+def run_yolo(img):
+
+    #img0 = Image.open(img.name).convert("RGB")
+    img_path = img.name # ["name"]
+    img0 = Image.open(img_path).convert("RGB")
+    draw = ImageDraw.Draw(img0)
+
+    predictions = predict(age_model_ts, model, stride, names, pt, jit, onnx, engine, imgsz=[320, 320], conf_thres=0.5, iou_thres=0.45, save_conf=True,
+                    exist_ok=True, nosave=True, save_txt=False, source=img_path, project=None, name=None)
+
+    detections : list[Detection] = []
+    for k, (bboxes, img)  in enumerate(predictions):
+        
+        #print(bboxes)
+        # exp.imgs.append(img_info)
+        for i, bbox in enumerate(bboxes):
+            det = Detection(
+                (k+1)*(i+1),
+                bbox["xmin"],
+                bbox["ymin"],
+                bbox["xmax"],
+                bbox["ymax"],
+                bbox["conf"],
+                bbox["class"],
+                bbox["class"],
+                img0.size
+            )
+            same = list(filter(lambda x: x.xmin == det.xmin and x.ymin == det.ymin or ( det.xmin > x.xmin and det.ymin > x.ymin and det.xmax < x.xmax and det.ymax < x.ymax ) or ( det.xmin < x.xmin and det.ymin < x.ymin and det.xmax > x.xmax and det.ymax > x.ymax ) or Detection.get_iou(det, x) > 0.6, detections))
+
+            if len(same) == 0:
+                detections.append(det)
+                draw.rectangle(((det.xmin, det.ymin), (det.xmax, det.ymax)), fill=None, outline=(255,255,255))
+                draw.rectangle(((det.xmin, det.ymin - 10), (det.xmax, det.ymin)), fill=(255,255,255))
+                draw.text((det.xmin, det.ymin - 10), det.class_name, fill=(0,0,0), font=ImageFont.truetype("Roboto-Regular.ttf"))
+
+    return img0
+
+inputs = gr.inputs.Image(type='file', label="Input Image")
+outputs = gr.outputs.Image(type="pil", label="Output Image")
+
+title = "AgeGuesser"
+description = "Guess the age of a person from his/her face!"
+article = """A fully automated system based on YOLOv5 and EfficientNet to perform face detection and age estimation in real-time. 
+    
+    Links:
+    <ul>
+    <li>
+    <a href='https://link.springer.com/chapter/10.1007/978-3-030-89131-2_25'>Paper</a>
+    </li>
+    <li>
+    <a href='https://www.researchgate.net/publication/355777953_Real-Time_Age_Estimation_from_Facial_Images_Using_YOLO_and_EfficientNet'>Paper</a> 
+    </li>
+    <li>
+    <a href='https://github.com/ai-hazard/AgeGuesser-train'>Github</a> 
+
+    </li>
+"""
+
+examples = [['images/1.jpg'], ['images/2.jpg'], ['images/3.jpg'], ['images/4.jpg'], ['images/5.jpg'], ]
+
+gr.Interface(run_yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, theme="huggingface").launch(enable_queue=True)