Views
No views yet
| Metric | Value |
|---|---|
| Input Values | 1,555,200 |
| Latent Values | 512 |
| Compression Ratio | 3037.50 : 1 |
| Data Retained | ~0.03% |
1import os
2import torch
3import torch.nn as nn
4from PIL import Image
5from torchvision import transforms
6import matplotlib.pyplot as plt
7
8# --- CONFIGURATION ---
9MODEL_REPO = "Parallax-labs-1/parallax-VISION_amongus-2.0"
10INPUT_IMAGE = "test_screenshot.png"
11DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
12LATENT_DIM = 512
13IMG_SIZE = 720
14
15# --- 1. MODEL ARCHITECTURE ---
16class ParallaxAutoencoder(nn.Module):
17 def __init__(self):
18 super(ParallaxAutoencoder, self).__init__()
19 # Encoder: Structural Downsampling
20 self.encoder = nn.Sequential(
21 nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
22 nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
23 nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
24 nn.Conv2d(64, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
25 nn.Flatten(),
26 nn.Linear(64 * 45 * 45, LATENT_DIM)
27 )
28 # Decoder: Geometric Expansion
29 self.decoder = nn.Sequential(
30 nn.Linear(LATENT_DIM, 64 * 45 * 45), nn.ReLU(),
31 nn.Unflatten(1, (64, 45, 45)),
32 nn.ConvTranspose2d(64, 64, 2, stride=2), nn.ReLU(),
33 nn.ConvTranspose2d(64, 32, 2, stride=2), nn.ReLU(),
34 nn.ConvTranspose2d(32, 16, 2, stride=2), nn.ReLU(),
35 nn.ConvTranspose2d(16, 3, 2, stride=2), nn.Sigmoid()
36 )
37
38 def forward(self, x):
39 return self.decoder(self.encoder(x))
40
41# --- 2. LOAD & INITIALIZE ---
42model = ParallaxAutoencoder().to(DEVICE)
43
44if not os.path.exists("model.pth"):
45 print("Downloading weights...")
46 os.system(f"wget https://huggingface.co/{MODEL_REPO}/resolve/main/model.pth")
47
48model.load_state_dict(torch.load('model.pth', map_location=DEVICE))
49model.eval()
50
51# --- 3. EXECUTE RECONSTRUCTION ---
52def run_reconstruction(img_path):
53 if not os.path.exists(img_path):
54 print(f"Error: {img_path} not found.")
55 return
56
57 img = Image.open(img_path).convert('RGB')
58 preprocess = transforms.Compose([
59 transforms.Resize((IMG_SIZE, IMG_SIZE)),
60 transforms.ToTensor()
61 ])
62 img_t = preprocess(img).unsqueeze(0).to(DEVICE)
63
64 with torch.no_grad():
65 output = model(img_t)
66
67 # Calculate Ratio
68 ratio = (IMG_SIZE * IMG_SIZE * 3) / LATENT_DIM
69
70 # Plotting
71 fig, ax = plt.subplots(1, 2, figsize=(15, 7))
72 ax[0].imshow(img_t[0].cpu().permute(1, 2, 0))
73 ax[0].set_title("Original Image")
74 ax[0].axis('off')
75
76 ax[1].imshow(output[0].cpu().permute(1, 2, 0))
77 ax[1].set_title(f"Reconstructed (Ratio {ratio:.2f}:1)")
78 ax[1].axis('off')
79
80 plt.show()
81
82if __name__ == "__main__":
83 run_reconstruction(INPUT_IMAGE)