Spaces:
Build error
Build error
File size: 2,899 Bytes
d7401b7 abb346a b361607 d7401b7 b361607 d7401b7 b361607 d7401b7 abb346a d7401b7 b361607 abb346a b361607 abb346a d7401b7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import gradio as gr
from ultralytics import YOLO
import os
from datetime import datetime
import requests
from requests.auth import HTTPBasicAuth
import base64
import cv2
from dotenv import load_dotenv
import io
model = YOLO('head_yolov8s.pt')
load_dotenv()
def upload_to_github(file_content, filepath):
username = "mohitsriv99"
repository = "hsamples"
branch = "main"
url = f"https://api.github.com/repos/{username}/{repository}/contents/{filepath}"
token = os.getenv("TOKEN") # Using an environment variable for the token
encoded_file = base64.b64encode(file_content).decode("utf-8")
payload = {
"message": f"Upload image {filepath}",
"content": encoded_file,
"branch": branch
}
headers = {
"Authorization": f"token {token}"
}
response = requests.put(url, json=payload, headers=headers)
if response.status_code == 201:
print("File uploaded successfully to GitHub.")
else:
print(f"Failed to upload file. Status code: {response.status_code}, Response: {response.text}")
def head_detect(path):
if path is None:
return None
img = cv2.imread(path)
# Generate a timestamped filename
current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"uploaded_images/{current_time}.jpg"
# Encode the image to a memory buffer instead of saving to disk
_, buffer = cv2.imencode('.jpg', img)
img_byte_array = io.BytesIO(buffer).getvalue()
# Upload the image to GitHub
upload_to_github(img_byte_array, filename)
output = model(source=img)
res = output[0].cpu().numpy()
# Extract bbox, cls id and conf
bboxes = res.boxes.xyxy
class_ids = res.boxes.cls
conf_scores = res.boxes.conf
for i in range(len(bboxes)):
xmin, ymin, xmax, ymax = int(bboxes[i][0]), int(bboxes[i][1]), int(bboxes[i][2]), int(bboxes[i][3])
conf = conf_scores[i]
cls_id = int(class_ids[i])
label = model.names[cls_id] # Get the label name
# Draw rectangle for bounding box
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color=(0, 0, 255), thickness=2, lineType=cv2.LINE_AA)
# Prepare label text with confidence score
label_text = f'{label} {conf:.2f}'
# Put text (label) on the image
cv2.putText(img, label_text, (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, lineType=cv2.LINE_AA)
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
input_image = [
gr.components.Image(type='filepath', label='Input Image'),
]
output_image = [
gr.components.Image(type='numpy', label='Prediction'),
]
interface = gr.Interface(
fn=head_detect,
inputs=input_image,
outputs=output_image,
title='Top Head Detection',
cache_examples=False,
)
gr.TabbedInterface(
[interface],
tab_names=['Image Inference']
).queue().launch() |