atticus-carter's picture
Update app.py
e37f9bc verified
Raw
History Blame Contribute Delete
1.53 kB
import os
import torch
# CRITICAL: Redirect cache to temporary storage to avoid Hugging Face Space eviction
os.environ['TORCH_HOME'] = '/tmp/torch_cache'
os.environ['HUB_DIR'] = '/tmp/torch_hub'
os.environ['TMPDIR'] = '/tmp'
# Added to fix the Ultralytics config warning:
os.environ['YOLO_CONFIG_DIR'] = '/tmp/yolo_config'
torch.hub.set_dir('/tmp/torch_hub')
import gradio as gr
from ultralytics import YOLO
# Load the model
model_path = "OceanCV_FirstPass.pt"
model = YOLO(model_path)
def run(image_path, conf, iou):
# Predict using the slider values
# Note: 'classes=0' is kept from your baseline template
results = model.predict(image_path, conf=conf, iou=iou, classes=0)
# Reverse channels from BGR (OpenCV/YOLO default) to RGB (Gradio expectation)
return results[0].plot()[:, :, ::-1]
title = "OceanCV First Pass Detector"
description = "Upload an image to detect objects using the first pass model."
# Build the interface with interactive sliders
interface = gr.Interface(
fn=run,
inputs=[
gr.Image(type="filepath", label="Upload Image"),
gr.Slider(minimum=0.05, maximum=1.0, value=0.20, step=0.05, label="Confidence Threshold"),
gr.Slider(minimum=0.05, maximum=1.0, value=0.85, step=0.05, label="IoU Threshold")
],
outputs=gr.Image(type="numpy", label="Detections"),
title=title,
description=description
)
# Launch the app with server_name and port explicitly set for HF Spaces
interface.queue().launch(server_name="0.0.0.0", server_port=7860)