LogicGoInfotechSpaces commited on
Commit
293fc40
·
verified ·
1 Parent(s): d8c6239

Update app/main.py

Browse files
Files changed (1) hide show
  1. app/main.py +7 -9
app/main.py CHANGED
@@ -35,10 +35,10 @@ except Exception as e:
35
  print("❌ Firebase initialization failed:", e)
36
 
37
  # -------------------------------------------------
38
- # 📁 Directories
39
  # -------------------------------------------------
40
- UPLOAD_DIR = "uploads"
41
- RESULTS_DIR = "results"
42
  os.makedirs(UPLOAD_DIR, exist_ok=True)
43
  os.makedirs(RESULTS_DIR, exist_ok=True)
44
 
@@ -54,16 +54,14 @@ model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
54
  print("📦 Loading model weights...")
55
  state_dict = torch.load(model_path, map_location="cpu")
56
 
57
- # NOTE:
58
- # Replace this with your actual model class.
59
- # This is just example placeholder.
60
  # from model import ColorizeNet
61
  # model = ColorizeNet()
62
  # model.load_state_dict(state_dict)
63
  # model.eval()
64
 
65
  def colorize_image(img: Image.Image):
66
- """ Dummy colorizer (convert grayscale fake color) """
67
  transform = transforms.ToTensor()
68
  tensor = transform(img.convert("L")).unsqueeze(0)
69
  tensor = tensor.repeat(1, 3, 1, 1)
@@ -98,7 +96,7 @@ async def upload_image(
98
  if not file.content_type.startswith("image/"):
99
  raise HTTPException(status_code=400, detail="Invalid file type")
100
 
101
- image_id = str(uuid.uuid4()) + ".jpg"
102
  file_path = os.path.join(UPLOAD_DIR, image_id)
103
 
104
  with open(file_path, "wb") as f:
@@ -128,7 +126,7 @@ async def colorize(
128
  img = Image.open(io.BytesIO(await file.read()))
129
  output_img = colorize_image(img)
130
 
131
- result_id = str(uuid.uuid4()) + ".jpg"
132
  output_path = os.path.join(RESULTS_DIR, result_id)
133
  output_img.save(output_path)
134
 
 
35
  print("❌ Firebase initialization failed:", e)
36
 
37
  # -------------------------------------------------
38
+ # 📁 Directories (FIXED FOR HUGGINGFACE SPACES)
39
  # -------------------------------------------------
40
+ UPLOAD_DIR = "/tmp/uploads"
41
+ RESULTS_DIR = "/tmp/results"
42
  os.makedirs(UPLOAD_DIR, exist_ok=True)
43
  os.makedirs(RESULTS_DIR, exist_ok=True)
44
 
 
54
  print("📦 Loading model weights...")
55
  state_dict = torch.load(model_path, map_location="cpu")
56
 
57
+ # NOTE: Replace with real model architecture
 
 
58
  # from model import ColorizeNet
59
  # model = ColorizeNet()
60
  # model.load_state_dict(state_dict)
61
  # model.eval()
62
 
63
  def colorize_image(img: Image.Image):
64
+ """ Dummy colorizer (replace with real model.predict) """
65
  transform = transforms.ToTensor()
66
  tensor = transform(img.convert("L")).unsqueeze(0)
67
  tensor = tensor.repeat(1, 3, 1, 1)
 
96
  if not file.content_type.startswith("image/"):
97
  raise HTTPException(status_code=400, detail="Invalid file type")
98
 
99
+ image_id = f"{uuid.uuid4()}.jpg"
100
  file_path = os.path.join(UPLOAD_DIR, image_id)
101
 
102
  with open(file_path, "wb") as f:
 
126
  img = Image.open(io.BytesIO(await file.read()))
127
  output_img = colorize_image(img)
128
 
129
+ result_id = f"{uuid.uuid4()}.jpg"
130
  output_path = os.path.join(RESULTS_DIR, result_id)
131
  output_img.save(output_path)
132