Fix launch-timeout crash and remove dead Dockerfile

Root cause: the yolov5 package's torch.hub download of the yolov5 repo
code, combined with unpinned CUDA torch wheels on cpu-basic hardware,
pushed the build well past the 30-minute health-check window. Redirect
Torch caches to /tmp, pin CPU-only torch/torchvision, patch torch.load
for PyTorch 2.6+'s weights_only default so the legacy yolov5 checkpoint
still loads, and drop the unused Dockerfile/tator_inference.py/tator
dependency left over from an old pipeline the gradio-SDK space never
actually used. Bump gradio to 6.9.0.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Dockerfile

@@ -1,10 +0,0 @@
-FROM python:3.7
-
-RUN apt-get update \
-    && apt-get install ffmpeg libsm6 libxext6 -y
-
-RUN pip install yolov5 tator gradio
-
-COPY . ./
-
-CMD [ "python", "-u", "./tator_inference.py" ]

README.md

@@ -4,7 +4,7 @@ emoji: 🐠
 colorFrom: indigo
 colorTo: gray
 sdk: gradio
-sdk_version: 4.36.1
+sdk_version: 6.9.0
 app_file: app.py
 pinned: false
 ---

inference.py

@@ -1,6 +1,25 @@
+import os
+
+# Keep Torch/YOLO hub caches off the Space's persistent disk so repeated
+# loads don't exceed the storage quota:
+# https://discuss.huggingface.co/t/how-to-fix-workload-evicted-storage-limit-exceeded-50g-error-in-huggingface-spaces/169258
+os.environ.setdefault("TORCH_HOME", "/tmp/torch_cache")
+os.environ.setdefault("HUB_DIR", "/tmp/torch_hub")
+os.environ.setdefault("TMPDIR", "/tmp")
+
 import glob
 import numpy as np
 import torch
+
+# PyTorch >=2.6 defaults torch.load(weights_only=True), which can't unpickle
+# this legacy yolov5 checkpoint format. We trust this checkpoint (it's our
+# own model), so restore the old default for it.
+_torch_load = torch.load
+def _torch_load_compat(*args, **kwargs):
+    kwargs.setdefault("weights_only", False)
+    return _torch_load(*args, **kwargs)
+torch.load = _torch_load_compat
+
 import yolov5
 from typing import Union, List, Optional
 

requirements.txt

@@ -1,3 +1,6 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.14.0
+torchvision==0.29.0
 numpy<2.0
 yolov5==6.2.3
-tator
+setuptools<81

tator_inference.py

@@ -1,100 +0,0 @@
-import os
-import logging
-from tempfile import TemporaryFile
-
-import cv2
-import numpy as np
-from PIL import Image
-
-import tator
-import inference
-
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-# Read environment variables that are provided from TATOR
-host = os.getenv('HOST')
-token = os.getenv('TOKEN')
-project_id = int(os.getenv('PROJECT_ID'))
-media_ids = [int(id_) for id_ in os.getenv('MEDIA_IDS').split(',')]
-frames_per_inference = int(os.getenv('FRAMES_PER_INFERENCE', 30))
-
-# Set up the TATOR API.
-api = tator.get_api(host, token)
-
-# Iterate through each video.
-for media_id in media_ids:
-
-    # Download video.
-    media = api.get_media(media_id)
-    logger.info(f"Downloading {media.name}...")
-    out_path = f"/tmp/{media.name}"
-    for progress in tator.util.download_media(api, media, out_path):
-        logger.info(f"Download progress: {progress}%")
-
-    # Do inference on each video.
-    logger.info(f"Doing inference on {media.name}...")
-    localizations = []
-    vid = cv2.VideoCapture(out_path)
-    frame_number = 0
-
-    # Read *every* frame from the video, break when at the end.
-    while True:
-        ret, frame = vid.read()
-        if not ret:
-            break
-
-        # Create a temporary file, access the image data, save data to file.
-        framefile = TemporaryFile(suffix='.jpg')
-        im = Image.fromarray(frame)
-        im.save(framefile)
-
-        # For every N frames, make a prediction; append prediction results
-        # to a list, increase the frame count.
-        if frame_number % frames_per_inference == 0:
-
-            spec = {}
-
-            # Predictions contains all information inside pandas dataframe
-            predictions = inference.run_inference(framefile)
-
-            for i, r in predictions.pandas().xyxy[0].iterrows:
-
-                spec['media_id'] = media_id
-                spec['type'] = None # Unsure, docs not specific
-                spec['frame'] = frame_number
-
-                x, y, x2, y2 = r['xmin'], r['ymin'], r['xmax'], r['ymax']
-                w, h = x2 - x, y2 - y
-
-                spec['x'] = x
-                spec['y'] = y
-                spec['width'] = w
-                spec['height'] = h
-                spec['class_category'] = r['name']
-                spec['confidence'] = r['confidence']
-
-                localizations.append(spec)
-
-        frame_number += 1
-
-    # End interaction with video properly.
-    vid.release()
-
-    logger.info(f"Uploading object detections on {media.name}...")
-
-    # Create the localizations in the video.
-    num_created = 0
-    for response in tator.util.chunked_create(api.create_localization_list,
-                                              project_id,
-                                              localization_spec=localizations):
-        num_created += len(response.id)
-
-    # Output pretty logging information.
-    logger.info(f"Successfully created {num_created} localizations on "
-                f"{media.name}!")
-
-    logger.info("-------------------------------------------------")
-
-logger.info(f"Completed inference on {len(media_ids)} files.")