File size: 3,225 Bytes
6a5bb7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """Background Difference Segmentation — Free Pixel-Level GT
=============================================================
Fixed cameras = static background. Compute per-camera background via
median of 100 images, then get goat masks via |frame - background|.
Output: Data/Detection_dataset/labels/train_seg/ with binary masks.
Runs on CPU.
"""
import sys,os
import numpy as np
from PIL import Image
from tqdm import tqdm
from collections import defaultdict
PROJECT_DIR='/home/user/goat'
os.chdir(PROJECT_DIR);sys.path.insert(0,PROJECT_DIR)
def compute_background(camera, img_dir, n_samples=100):
"""Compute median background for one camera."""
files=[f for f in os.listdir(img_dir) if camera in f and f.endswith('.jpg')]
files=sorted(files)[:n_samples]
if not files: return None
# Load all images
imgs=[]
for f in tqdm(files, desc=f'{camera} bg load'):
img=np.array(Image.open(os.path.join(img_dir,f)).convert('L'),dtype=np.float32)
imgs.append(img)
# Median per pixel
stack=np.stack(imgs,axis=0)
bg=np.median(stack,axis=0).astype(np.uint8)
return bg
def main():
img_dir='Data/Detection_dataset/images/train'
cameras=['EastLeft','EastRight','WestLeft','WestRight']
# Compute backgrounds
backgrounds={}
for cam in cameras:
bg=compute_background(cam, img_dir)
if bg is not None:
backgrounds[cam]=bg
# Save
out_path=f'Data/background_{cam}.png'
Image.fromarray(bg).save(out_path)
print(f'{cam} background saved: {out_path} ({bg.shape})')
# Generate masks for all training images
out_dir='Data/Detection_dataset/labels/train_seg'
os.makedirs(out_dir,exist_ok=True)
all_files=sorted([f for f in os.listdir(img_dir) if f.endswith('.jpg')])
total_masks=0
for f in tqdm(all_files, desc='Generating masks'):
# Determine camera
cam=None
for c in cameras:
if c in f: cam=c;break
if cam is None or cam not in backgrounds: continue
img=np.array(Image.open(os.path.join(img_dir,f)).convert('L'),dtype=np.float32)
bg=backgrounds[cam].astype(np.float32)
# |frame - background|
diff=np.abs(img-bg)
# Adaptive threshold: mean + 2*std of difference
thresh=diff.mean()+2.0*diff.std()
mask=(diff>thresh).astype(np.uint8)*255
# Clean up: remove small noise
from scipy import ndimage
mask=ndimage.binary_fill_holes(mask).astype(np.uint8)*255
# Remove objects smaller than 100 pixels
labeled,_=ndimage.label(mask)
sizes=ndimage.sum(mask,labeled,range(1,labeled.max()+1))
clean_mask=np.zeros_like(mask)
for i,sz in enumerate(sizes,1):
if sz>100:
clean_mask[labeled==i]=255
mask=clean_mask
# Save
out_name=f.replace('.jpg','.png')
Image.fromarray(mask).save(os.path.join(out_dir,out_name))
# Count pixels
goat_pixels=(mask>0).sum()
total_masks+=1
print(f'\nSegmentation masks generated: {total_masks} images')
print(f'Saved to: {out_dir}')
return out_dir
if __name__=='__main__':
main()
|