Skip to content
Snippets Groups Projects
Commit 2e77c4f5 authored by fresleven's avatar fresleven
Browse files

update simdata

parent ede8990e
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 219 additions and 0 deletions
crop/SimData.ipynb 0 → 100644
+ 219
− 0
View file @ 2e77c4f5
%% Cell type:code id: tags:
``` python
import glob
from PIL import Image
import pillow_heif
import pandas as pd
import numpy as np
from torchvision.utils import save_image
import torchvision.transforms.functional as fn
```
%% Cell type:code id: tags:
``` python
# TODO: overlap?
def generate_data(bg: list[Image], beetle_set: list[Image], num_beetles_arr: list[int]):
set_size = len(beetle_set)
set_bg_size = len(bg)
sim_arr, coords_arr = [],[]
for num_beetles in num_beetles_arr:
bg_id = np.random.randint(0, set_bg_size)
bg_temp = bg[bg_id].copy()
width, height = bg_temp.size
beetle_coords = []
for _ in range(num_beetles):
# get random beetle image
beetle_id = np.random.randint(0, set_size)
beetle_img = beetle_set[beetle_id]
beetle_width, beetle_height = beetle_img.size
# get random x,y coords to paste beetle
x = randint(0, width - beetle_width)
y = randint(0, height - beetle_height)
# get random beetle rotation
angle = randint(0, 359)
beetle_img = beetle_img.rotate(angle, resample=Image.BICUBIC)
bg_temp.paste(beetle_img, box=(x,y), mask=beetle_img)
beetle_coords.append((beetle_id, x, y, beetle_width, beetle_height, angle))
sim_arr.append(bg_temp)
coords_arr.append(beetle_coords)
return sim_arr, coords_arr
```
%% Cell type:code id: tags:
``` python
def find_coeffs(source_coords, target_coords):
matrix = []
for s, t in zip(source_coords, target_coords):
matrix.append([t[0], t[1], 1, 0, 0, 0, -s[0]*t[0], -s[0]*t[1]])
matrix.append([0, 0, 0, t[0], t[1], 1, -s[1]*t[0], -s[1]*t[1]])
A = np.matrix(matrix, dtype=float)
B = np.array(source_coords).reshape(8)
res = np.dot(np.linalg.inv(A.T * A) * A.T, B)
return np.array(res).reshape(8)
```
%% Cell type:code id: tags:
``` python
bg = Image.open("bg.png")
beetles = []
for file in glob.glob(r"/raid/projects/akhot2/group-01-phys371-sp2023/crop/beetles/*"):
b0 = Image.open(file)
beetles.append(b0);
# map corners of trap to corners of image
coeffs = find_coeffs([(128,6), (1904,62), (2113,3137), (3,3228)],
[(0,0), (bg.size[0], 0), (bg.size[0], bg.size[1]), (0, bg.size[1])])
bg_flat = bg.transform(bg.size, Image.PERSPECTIVE, coeffs, Image.BICUBIC)
```
%% Cell type:code id: tags:
``` python
beetle_counts = np.random.randint(0, 6, size=1)
sim_img_arr, coords_arr = generate_data([bg_flat], beetles, beetle_counts)
```
%% Cell type:code id: tags:
``` python
def export(img_arr, coords_arr):
transform = transforms.Compose([transforms.ToTensor()])
s = ""
image_set = "train"
change = len(img_arr)*0.8
for i in range(len(img_arr)):
s = ""
for coord in coords_arr[i]:
c, x, y, w, h, a = coord
img = img_arr[i]
s += f"{0} {x/img.size[0]} {y/img.size[1]} {w/img.size[0]} {h/img.size[1]}\n"
if i > change:
image_set = "test"
with open("data/" +image_set+ f"/labels/sim{i}.txt", "w") as f:
f.write(s)
save_image(img, "data/" +image_set+ f"/images/sim{i}.png")
```
%% Cell type:code id: tags:
``` python
export(sim_img_arr, coords_arr)
```
%% Output
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Input In [51], in <cell line: 1>()
----> 1 export(sim_img_arr, coords_arr)
Input In [50], in export(img_arr, coords_arr)
13 with open("data/" +image_set+ f"/labels/sim{i}.txt", "w") as f:
14 f.write(s)
---> 15 save_image(img, "data/" +image_set+ f"/images/sim{i}.png")
File /raid/projects/akhot2/conda/envs/akhot2/lib/python3.9/site-packages/torch/autograd/grad_mode.py:27, in _DecoratorContextManager.__call__.<locals>.decorate_context(*args, **kwargs)
24 @functools.wraps(func)
25 def decorate_context(*args, **kwargs):
26 with self.clone():
---> 27 return func(*args, **kwargs)
File /raid/projects/akhot2/conda/envs/akhot2/lib/python3.9/site-packages/torchvision/utils.py:152, in save_image(tensor, fp, format, **kwargs)
150 if not torch.jit.is_scripting() and not torch.jit.is_tracing():
151 _log_api_usage_once(save_image)
--> 152 grid = make_grid(tensor, **kwargs)
153 # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
154 ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
File /raid/projects/akhot2/conda/envs/akhot2/lib/python3.9/site-packages/torch/autograd/grad_mode.py:27, in _DecoratorContextManager.__call__.<locals>.decorate_context(*args, **kwargs)
24 @functools.wraps(func)
25 def decorate_context(*args, **kwargs):
26 with self.clone():
---> 27 return func(*args, **kwargs)
File /raid/projects/akhot2/conda/envs/akhot2/lib/python3.9/site-packages/torchvision/utils.py:60, in make_grid(tensor, nrow, padding, normalize, value_range, scale_each, pad_value, **kwargs)
58 _log_api_usage_once(make_grid)
59 if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
---> 60 raise TypeError(f"tensor or list of tensors expected, got {type(tensor)}")
62 if "range" in kwargs.keys():
63 warnings.warn(
64 "The parameter 'range' is deprecated since 0.12 and will be removed in 0.14. "
65 "Please use 'value_range' instead."
66 )
TypeError: tensor or list of tensors expected, got <class 'PIL.Image.Image'>
%% Cell type:code id: tags:
``` python
format
0 x y width height (each value normalized to size of image)
```
%% Cell type:code id: tags:
``` python
len(beetle_counts)
```
%% Output
100
%% Cell type:code id: tags:
``` python
```
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment