Update app.py

app.py

@@ -1,136 +1,42 @@
+from flask import Flask, request, jsonify, send_file
+from flask_cors import CORS
 import os
-import torch
-import gradio as gr
-import numpy as np
 from PIL import Image
-import tempfile
-from skimage import measure
-import trimesh
-import torch.nn.functional as F
-import torchvision.transforms as transforms
+import io
+import torch
+from gfpgan import GFPGANer
 
-# Check if CUDA is available, otherwise use CPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+app = Flask(__name__)
+CORS(app)
 
-# Define a simple neural network to extract depth from images
-class SimpleDepthNet(torch.nn.Module):
-    def __init__(self):
-        super(SimpleDepthNet, self).__init__()
-        self.conv1 = torch.nn.Conv2d(3, 32, kernel_size=3, padding=1)
-        self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=3, padding=1)
-        self.conv3 = torch.nn.Conv2d(64, 128, kernel_size=3, padding=1)
-        self.conv4 = torch.nn.Conv2d(128, 1, kernel_size=3, padding=1)
-        self.pool = torch.nn.MaxPool2d(2, 2)
-        self.upsample = torch.nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
-        
-    def forward(self, x):
-        # Encoder
-        x = F.relu(self.conv1(x))
-        x = self.pool(x)
-        x = F.relu(self.conv2(x))
-        x = self.pool(x)
-        
-        # Decoder
-        x = self.upsample(x)
-        x = F.relu(self.conv3(x))
-        x = self.upsample(x)
-        x = torch.sigmoid(self.conv4(x))
-        return x
+model = GFPGANer(
+    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.3.pth',
+    upscale=2,
+    arch='clean',
+    channel_multiplier=2,
+    bg_upsampler=None,
+    device='cpu'
+)
 
-# Initialize the model
-model = SimpleDepthNet().to(device)
+@app.route('/enhance', methods=['POST'])
+def enhance():
+    if 'image' not in request.files:
+        return jsonify({'error': 'No image uploaded'}), 400
 
-# Define transformation for input images
-transform = transforms.Compose([
-    transforms.Resize((256, 256)),
-    transforms.ToTensor(),
-])
+    img_file = request.files['image']
+    img = Image.open(img_file.stream).convert('RGB')
 
-def image_to_3d(image):
-    """
-    Convert a single image to a 3D model using a simple depth extraction approach
-    """
-    if image is None:
-        return None, "No image provided"
-    
     try:
-        # Preprocess image
-        img_tensor = transform(image).unsqueeze(0).to(device)
-        
-        # Generate depth map
-        with torch.no_grad():
-            depth = model(img_tensor)[0, 0].cpu().numpy()
-        
-        # Convert depth map to 3D points
-        h, w = depth.shape
-        y, x = np.meshgrid(np.arange(h), np.arange(w), indexing='ij')
-        
-        # Normalize coordinates
-        x = (x - w/2) / max(w, h)
-        y = (y - h/2) / max(w, h)
-        z = depth - 0.5  # Center around zero
-        
-        # Create point cloud
-        points = np.stack([x.flatten(), y.flatten(), z.flatten()], axis=1)
-        
-        # Get colors from original image
-        img_np = np.array(image.resize((w, h))) / 255.0
-        colors = img_np.reshape(-1, 3)
-        
-        # Create a mesh from the point cloud (using marching cubes on the depth map)
-        verts, faces, _, _ = measure.marching_cubes(depth, 0.5)
-        mesh = trimesh.Trimesh(vertices=verts, faces=faces)
-        
-        # Save as OBJ
-        with tempfile.NamedTemporaryFile(suffix='.obj', delete=False) as obj_file:
-            obj_path = obj_file.name
-            mesh.export(obj_path)
-        
-        # Also save as PLY for better compatibility with Unity
-        with tempfile.NamedTemporaryFile(suffix='.ply', delete=False) as ply_file:
-            ply_path = ply_file.name
-            mesh.export(ply_path)
-        
-        return [obj_path, ply_path], "3D model generated successfully!"
-    except Exception as e:
-        return None, f"Error: {str(e)}"
+        _, _, restored_img = model.enhance(np.array(img), has_aligned=False, only_center_face=True, paste_back=True)
+        result_img = Image.fromarray(restored_img)
 
-def process_image(image):
-    try:
-        if image is None:
-            return None, None, "Please upload an image first."
-        
-        results, message = image_to_3d(image)
-        
-        if results:
-            return results[0], results[1], message
-        else:
-            return None, None, message
-    except Exception as e:
-        return None, None, f"Error: {str(e)}"
+        img_io = io.BytesIO()
+        result_img.save(img_io, 'JPEG')
+        img_io.seek(0)
 
-# Create Gradio interface
-with gr.Blocks(title="Simple Image to 3D Converter") as demo:
-    gr.Markdown("# Simple Image to 3D Converter")
-    gr.Markdown("Upload an image to convert it to a simple 3D model that you can use in Unity or other engines.")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            input_image = gr.Image(type="pil", label="Input Image")
-            submit_btn = gr.Button("Convert to 3D")
-        
-        with gr.Column(scale=1):
-            obj_file = gr.File(label="OBJ File (for editing)")
-            ply_file = gr.File(label="PLY File (for Unity)")
-            output_message = gr.Textbox(label="Output Message")
-    
-    submit_btn.click(
-        fn=process_image,
-        inputs=[input_image],
-        outputs=[obj_file, ply_file, output_message]
-    )
+        return send_file(img_io, mimetype='image/jpeg')
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
 
-# Launch the app
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000)