Update app.py

app.py

@@ -44,6 +44,9 @@ transform = transforms.Compose([GOSNormalize([0.5,0.5,0.5],[1.0,1.0,1.0])])
 
 def load_image(im_path, hypar):
     im = im_reader(im_path)
+    # Convert to RGB if image has alpha channel
+    if im.mode == 'RGBA':
+        im = im.convert('RGB')
     im, im_shp = im_preprocess(im, hypar["cache_size"])
     im = torch.divide(im,255.0)
     shape = torch.from_numpy(np.array(im_shp))
@@ -153,7 +156,8 @@ def remove_background():
         
         return jsonify({
             'rgba_image': f"/results/rgba_{file.filename}",
-            'mask_image': f"/results/mask_{file.filename}"
+            'mask_image': f"/results/mask_{file.filename}",
+            'original_filename': file.filename
         })
     except Exception as e:
         return jsonify({'error': str(e)}), 500